{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e1e2e446",
   "metadata": {},
   "source": [
    "This is part 1 of a tutorial series. We recommend following them in order, starting with [Part 0: Welcome to `musica`](0.%20Welcome%20to%20MUSICA.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7030540f",
   "metadata": {},
   "source": [
    "# Multiple Grid Cells in `musica`\n",
    "\n",
    "As we saw in the last tutorial, gas-phase systems can be solved in `musica` using a `solver` and a `state`. For box models, that's all we need. For more complex applications (columns, 3D grids, etc.), we'll need to manage the state of a number of independent well-mixed air masses, or \"grid cells\".\n",
    "\n",
    ">__NOTE:__ We use \"grid cell\" throughout these tutorials to mean a well-mixed air mass. A vertical stack of \"grid cells\" would be used for a column-model. Collections of vertical stacks of grid cells would be used for a 3D grid or mesh. However, `micm` does not make assumptions about the shape, size, or arragement of the well-mixed air mass(es) a `state` represents, so the state can be used to represent air masses of any size/shape in any arrangement.\n",
    "\n",
    "_Why multiple grid cell states?_\n",
    "\n",
    "If we can create as many states as we want and use the same solver on them, maybe you're wondering why we even need multiple-grid-cell states. The answer, once again, is for performance. In HPC applications, `micm` can be used to solve hundreds of thousands of independent grid cells simultaneously. To do this efficiently, the data needs to be efficiently arranaged in memory. Multiple-grid-cell states allow us to do this."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "008e870b",
   "metadata": {},
   "source": [
    "## 1. Importing Libraries\n",
    "Below is a list of the required libraries for this tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c921c61",
   "metadata": {},
   "outputs": [],
   "source": [
    "import musica\n",
    "import musica.mechanism_configuration as mc\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "pd.set_option('display.float_format', str) # This is done to make the arrays more readable\n",
    "np.set_printoptions(suppress=True) # This is done to make the arrays more readable"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01472a4f",
   "metadata": {},
   "source": [
    "## 2. Defining a System\n",
    "\n",
    "In the first tutorial, we used an example configuration file to define our chemical system. This time, we'll use the `musica` API to define a chemical system in code. We'll use this apporach for future tutorials as well.\n",
    "\n",
    "Let's set up a system analogous to the simple analytically solvable system we used in the first tutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b1084ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = mc.Species(name=\"A\")\n",
    "B = mc.Species(name=\"B\")\n",
    "C = mc.Species(name=\"C\")\n",
    "species = [A, B, C]\n",
    "gas = mc.Phase(name=\"gas\", species=species)\n",
    "\n",
    "r1 = mc.Arrhenius(\n",
    "    name=\"A_to_B\",\n",
    "    A=4.0e-3,\n",
    "    C=50,\n",
    "    reactants=[A],\n",
    "    products=[B],\n",
    "    gas_phase=gas\n",
    ")\n",
    "\n",
    "r2 = mc.Arrhenius(\n",
    "    name=\"B_to_C\",\n",
    "    A=4.0e-3,\n",
    "    C=50,  \n",
    "    reactants=[B],\n",
    "    products=[C],\n",
    "    gas_phase=gas\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45b4a60a",
   "metadata": {},
   "source": [
    "Similar to the last tutorial, we have defined a simple system with:\n",
    "* Three species (`A`, `B`, and `C`) in the gas-phase\n",
    "* Two Arrhenius reactions\n",
    "    * `A -> B`\n",
    "    * `B -> C`\n",
    "\n",
    "We try to keep the python API consistent with the format of the configuration files, and have both structured around the science.\n",
    "\n",
    "The documentation can be used to clarify important details, like the parameter names and units for rate constant types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81bcef4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mc.Arrhenius?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ccc93f7",
   "metadata": {},
   "source": [
    "In the last chapter, we used the `musica` parser to turn our configuration file into a `mechanism`. Here, we'll create the mechanism out of the chemistry objects we've defined above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80008904-7666-449e-a527-8058e5194512",
   "metadata": {},
   "outputs": [],
   "source": [
    "mechanism = mc.Mechanism(\n",
    "    name=\"musica_micm_example\",\n",
    "    species=species,\n",
    "    phases=[gas],\n",
    "    reactions=[r1, r2]\n",
    ")\n",
    "for reaction in mechanism.reactions:\n",
    "    print(f\"[{reaction.type.name}] {reaction.to_equation()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7861e510",
   "metadata": {},
   "source": [
    "## 3. Creating the Solver\n",
    "\n",
    "We create a solver for our `mechanism` the same way we did in the last tutorial.\n",
    "One minor difference here, is that we specifcy the `solver_type`.\n",
    "For more information on the types of solvers available, see [here](https://ncar.github.io/micm/user_guide/solver_configurations.html).\n",
    "The standard-ordered Rosenbrock solver is actually the default `solver_type`, so including it here has no real effect.\n",
    "\n",
    ">__NOTE:__ If you're wondering what \"standard order\" means, it refers to how the matrix data in the `state` in arranged in memory. `micm` is designed to be performant enough to be included in large-scale climate and weather models. To achieve optimal performance we have advanced matrix ordering strategies we can apply to encourage vectorization of the solver instructions. For most Python-based applications, the standard ordering will be sufficient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18bb7d49",
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = musica.MICM(mechanism = mechanism, solver_type = musica.SolverType.rosenbrock_standard_order)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2199a26f",
   "metadata": {},
   "source": [
    "## 4. Creating the State\n",
    "\n",
    "In the last tutorial, we created a single-grid-cell `state`. This time, let's double the complexity of our system and create a two-grid-cell `state`!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d91a4be8",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_grid_cells = 2\n",
    "state = solver.create_state(num_grid_cells)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08d6e59c",
   "metadata": {},
   "source": [
    "## 5. Setting Conditions\n",
    "\n",
    "Next, let's set the conditions for our two-grid-cell state. In the last tutorial, you may have wondered why the temperature, pressure, and species concentrations were in array format. Hopefully, now it's clear - one element for each grid cell!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0a74a68",
   "metadata": {},
   "outputs": [],
   "source": [
    "state.set_conditions(temperatures=[300, 100], pressures=[101253.3, 11253.3])\n",
    "state.set_concentrations({\n",
    "    'A': [5, 20],\n",
    "    'B': [5, 3],\n",
    "    'C': [5, 7]\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1238f02b",
   "metadata": {},
   "source": [
    "## 7. Solving the System\n",
    "\n",
    "This time around, we'll advance the state for 60 s, collecting intermediate states each second."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76b21308",
   "metadata": {},
   "outputs": [],
   "source": [
    "concentrations_solved = []\n",
    "time_step = 10 # s\n",
    "sim_length = 600 # s\n",
    "curr_time = 0 # s\n",
    "while curr_time <= sim_length:\n",
    "    solver.solve(state, time_step)\n",
    "    concentrations_solved.append(state.get_concentrations())\n",
    "    curr_time += time_step"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d2185c0",
   "metadata": {},
   "source": [
    "## 8. Preparing the Results (Advanced; Optional Read)\n",
    "\n",
    "We're going to create a short helper function to put the results into a Pandas DataFrame for easier plotting.\n",
    "The function will extract results for a single grid cell, and we call it twice, once for each of the grid cells in our `state`.\n",
    "In this simulation, the temperature, pressure, and air density are all constant, so numpy's `repeat()` function is used to repeat their respective values for every time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65de1f01",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_results_single_cell(cell_index):\n",
    "    concentrations_solved_pd = []\n",
    "    for i in range(0, sim_length + 1, time_step):\n",
    "        concentrations_solved_pd.append({species: concentration[cell_index] for species, concentration in concentrations_solved[int(i/time_step)].items()})\n",
    "    df = pd.DataFrame(concentrations_solved_pd)\n",
    "    df = df.rename(columns = {'A' : 'CONC.A.mol m-3', 'B' : 'CONC.B.mol m-3', 'C' : 'CONC.C.mol m-3'})\n",
    "    df['time.s'] = list(map(float, range(0, sim_length + 1, time_step)))\n",
    "    df['ENV.temperature.K'] = np.repeat(state.get_conditions()['temperature'][cell_index], sim_length/time_step + 1.0)\n",
    "    df['ENV.pressure.Pa'] = np.repeat(state.get_conditions()['pressure'][cell_index], sim_length/time_step + 1.0)\n",
    "    df['ENV.air number density.mol m-3'] = np.repeat(state.get_conditions()['air_density'][cell_index], sim_length/time_step + 1.0)\n",
    "    df = df[['time.s', 'ENV.temperature.K', 'ENV.pressure.Pa', 'ENV.air number density.mol m-3', 'CONC.A.mol m-3', 'CONC.B.mol m-3', 'CONC.C.mol m-3']]\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9dda1ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_0 = convert_results_single_cell(0)\n",
    "df_1 = convert_results_single_cell(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ada2398b",
   "metadata": {},
   "source": [
    "## 9. Viewing the Results\n",
    "\n",
    "Finally, let's desplay and plot the DataFrames to show the evolution of both of the independent systems over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d49f1b08",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(df_0)\n",
    "display(df_1)\n",
    "df_0.plot(x='time.s', y=['CONC.A.mol m-3', 'CONC.B.mol m-3', 'CONC.C.mol m-3'], title='Concentration over time (grid cell 0)', ylabel='Concentration (mol m-3)', xlabel='Time (s)')\n",
    "df_1.plot(x='time.s', y=['CONC.A.mol m-3', 'CONC.B.mol m-3', 'CONC.C.mol m-3'], title='Concentration over time (grid cell 1)', ylabel='Concentration (mol m-3)', xlabel='Time (s)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dbb6055",
   "metadata": {},
   "source": [
    "## 10. Going Full Circle\n",
    "\n",
    "We've now seen the two ways to configure `musica` in practice: using configuration files, and in-code using the `musica` API. What if we want to create a configuration code from our in-code mechanism?\n",
    "\n",
    "Let's give it a try!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d8a0f5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "mechanism.export(\"my_cool_musica_mechanism.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9553e185",
   "metadata": {},
   "source": [
    "Let's take a look at what it looks like"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fb00e50",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open(\"my_cool_musica_mechanism.json\", \"r\") as f:\n",
    "    config = json.load(f)\n",
    "\n",
    "print(json.dumps(config, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63507e4d",
   "metadata": {},
   "source": [
    "Why is this useful? The `musica` library than underpins the `musica` python package is also used in large 3D models. For those applications, text-based configuration files allow run-time specification of the chemistry system without the need to modify source code and recompile the models. So, you can develop and test your mechanism in Python and when you're ready, turn it into a configuration file that you can use in a model like CheMPAS-A or CAM-SIMA!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.12.3)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}